# -*- coding: utf-8 -*-
"""
Created on Mon Oct  8 14:32:52 2018

@author: luolei

粗样本共线性分析
"""
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
import numpy as np
import sys

sys.path.append('../')

from mod import *


def cal_vif(y_true, y_pred):
	"""计算方差膨胀因子， VIF = 1 / (1 - R^2)"""
	y_true = y_true.flatten()
	y_pred = y_pred.flatten()
	
	r2 = r2_score(y_true, y_pred)
	vif = 1 / (1 - r2)
	return vif


def multicollinearity_filtering(total_samples, vif_thres = 10, seg_len = 10000):
	"""通过方差膨胀因子VIF进行特征筛选"""
	chosen_cols = []
	for col in selected_cols:
		col_names = [p for p in total_samples.columns if (col in p) & (p[len(col):].count('_') == 1)]
		col_n = col_names.__len__()
		col_samples = total_samples.loc[:seg_len, ['time'] + col_names].copy()  # **只用一段长度的样本计算VIF
		
		chosen_ids = [0]
		samples_arr = col_samples.loc[:, col + '_{}'.format(0)].to_numpy().reshape(-1, 1)
		for i in range(1, col_n):
			# 候选特征
			candidate_arr = col_samples.loc[:, col + '_{}'.format(i)].to_numpy().reshape(-1, 1)
			
			# 建立线性回归模型，计算VIF值
			rgsr = LinearRegression()
			rgsr.fit(samples_arr, candidate_arr)
			
			y_pred = rgsr.predict(samples_arr)
			vif = cal_vif(candidate_arr, y_pred)
			
			# 根据VIF结果判断共线性程度
			if vif < vif_thres:
				print('choose {}_{}'.format(col, i))
				samples_arr = np.hstack((samples_arr, candidate_arr))
				chosen_ids.append(i)
			else:
				# print('drop {}_{}'.format(col, i))
				pass
		
		chosen_cols += [col + '_{}'.format(p) for p in chosen_ids]
	
	total_samples = total_samples.loc[:, ['time'] + chosen_cols]
	return total_samples
		

if __name__ == '__main__':
	#%% 载入数据，构造样本
	import pandas as pd
	from mod.modeling.build_samples import build_samples_dataframe
	data = pd.read_csv('../../data/runtime/cstr_data.csv')
	total_samples, total_cols_n, cols, embed_dims = build_samples_dataframe(data)